import pickle
import _pickle
import os
import numpy as np

salBasePath = "D:/VR_project/PanoSaliency/data/"
saveSalPath = "D:/VR_project/ViewPrediction/frames/attentionSaliency/"
saveTimePath = "D:/VR_project/ViewPrediction/frames/attentionTimeStamp/"

fileDict = {0: 'Conan1', 1: 'Skiing', 2: 'Alien', 3: 'Conan2', 4: 'Surfing',
            5: 'War', 6: 'Cooking', 7: 'Football', 8: 'Rhinos'}
salFileList = ["saliency_ds2_topic0", "saliency_ds2_topic1", "saliency_ds2_topic2",
               "saliency_ds2_topic3", "saliency_ds2_topic4", "saliency_ds2_topic5",
               "saliency_ds2_topic6", "saliency_ds2_topic7", "saliency_ds2_topic8"]

for i in range(len(salFileList)):
    saliencyPath = salBasePath + salFileList[i]
    try:
        saliency_array = np.array(pickle.load(open(saliencyPath, 'rb'), encoding='bytes'), dtype=object)
    except _pickle.UnpicklingError:
        saliency_array = np.load(saliencyPath, allow_pickle=True)

    time_list = [item[0] for item in saliency_array]
    saliency_list = [item[2] for item in saliency_array]

    sample_time_list = time_list[::4]
    sample_saliency_list = saliency_list[::4]

    np.savetxt(saveTimePath + f"{i}-{fileDict[i]}.txt", np.array(sample_time_list), fmt='%.3f')
    np.save(saveSalPath + f"{i}-{fileDict[i]}.npy", np.array(sample_saliency_list))

    print(f"sample {salFileList[i]} finish!")